import socket
import websockets
import asyncio
import struct


# -------------------------- 1. 产品接口（抽象产品） --------------------------
# 连接接口
class NetworkConnection:
    def connect(self) -> bool:
        raise NotImplementedError("子类必须实现connect方法")

    def send_data(self, data: bytes) -> bool:
        raise NotImplementedError("子类必须实现send_data方法")

    def close(self) -> None:
        raise NotImplementedError("子类必须实现close方法")


# 数据处理器接口（与连接配套的产品）
class DataProcessor:
    def encode(self, data: bytes) -> bytes:
        """编码数据（如添加长度前缀、校验和等）"""
        raise NotImplementedError("子类必须实现encode方法")

    def decode(self, data: bytes) -> bytes:
        """解码数据"""
        raise NotImplementedError("子类必须实现decode方法")


# -------------------------- 2. TCP产品族 --------------------------
class TcpConnection(NetworkConnection):
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.socket = None

    def connect(self) -> bool:
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((self.host, self.port))
            print(f"[TCP] 连接成功：{self.host}:{self.port}")
            return True
        except Exception as e:
            print(f"[TCP] 连接失败：{e}")
            return False

    def send_data(self, data: bytes) -> bool:
        if not self.socket:
            print("[TCP] 连接未建立，无法发送")
            return False
        try:
            self.socket.sendall(data)
            print(f"[TCP] 发送成功：{len(data)}字节")
            return True
        except Exception as e:
            print(f"[TCP] 发送失败：{e}")
            return False

    def close(self) -> None:
        if self.socket:
            self.socket.close()
            print(f"[TCP] 连接关闭：{self.host}:{self.port}")
            self.socket = None


class TcpDataProcessor(DataProcessor):
    """TCP数据处理器（处理粘包：添加4字节长度前缀）"""

    def encode(self, data: bytes) -> bytes:
        # 前4字节为数据长度（大端模式）
        length = struct.pack('>I', len(data))
        return length + data

    def decode(self, data: bytes) -> bytes:
        # 移除长度前缀
        if len(data) < 4:
            return b""
        return data[4:]


# -------------------------- 3. UDP产品族 --------------------------
class UdpConnection(NetworkConnection):
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.socket = None

    def connect(self) -> bool:
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            print(f"[UDP] 初始化成功（目标：{self.host}:{self.port}）")
            return True
        except Exception as e:
            print(f"[UDP] 初始化失败：{e}")
            return False

    def send_data(self, data: bytes) -> bool:
        if not self.socket:
            print("[UDP] 未初始化，无法发送")
            return False
        try:
            self.socket.sendto(data, (self.host, self.port))
            print(f"[UDP] 发送成功：{len(data)}字节")
            return True
        except Exception as e:
            print(f"[UDP] 发送失败：{e}")
            return False

    def close(self) -> None:
        if self.socket:
            self.socket.close()
            print("[UDP] 连接关闭")
            self.socket = None


class UdpDataProcessor(DataProcessor):
    """UDP数据处理器（添加简单校验和）"""

    def encode(self, data: bytes) -> bytes:
        # 最后1字节为校验和（简单异或）
        checksum = 0
        for b in data:
            checksum ^= b
        return data + bytes([checksum])

    def decode(self, data: bytes) -> bytes:
        # 验证并移除校验和
        if len(data) < 1:
            return b""
        checksum = 0
        for b in data[:-1]:
            checksum ^= b
        if checksum == data[-1]:
            return data[:-1]  # 校验通过
        return b""  # 校验失败返回空


# -------------------------- 4. WebSocket产品族 --------------------------
class WsConnection(NetworkConnection):
    def __init__(self, url: str):
        self.url = url
        self.websocket = None
        self.loop = None

    def connect(self) -> bool:
        try:
            self.loop = asyncio.get_event_loop()
            self.websocket = self.loop.run_until_complete(websockets.connect(self.url))
            print(f"[WebSocket] 连接成功：{self.url}")
            return True
        except Exception as e:
            print(f"[WebSocket] 连接失败：{e}")
            return False

    def send_data(self, data: bytes) -> bool:
        if not self.websocket or not self.loop:
            print("[WebSocket] 连接未建立，无法发送")
            return False
        try:
            self.loop.run_until_complete(self.websocket.send(data))
            print(f"[WebSocket] 发送成功：{len(data)}字节")
            return True
        except Exception as e:
            print(f"[WebSocket] 发送失败：{e}")
            return False

    def close(self) -> None:
        if self.websocket and self.loop:
            self.loop.run_until_complete(self.websocket.close())
            print(f"[WebSocket] 连接关闭：{self.url}")
            self.websocket = None
            self.loop = None


class WsDataProcessor(DataProcessor):
    """WebSocket数据处理器（JSON格式化）"""

    def encode(self, data: bytes) -> bytes:
        # 简单JSON包装
        return f'{{"type":"data","content":"{data.decode()}","length":{len(data)}}}\n'.encode()

    def decode(self, data: bytes) -> bytes:
        # 提取content字段
        import json
        try:
            obj = json.loads(data.decode())
            return obj.get("content", "").encode()
        except:
            return b""


# -------------------------- 5. 抽象工厂接口 --------------------------
class AbstractNetworkFactory:
    def create_connection(self, **kwargs) -> NetworkConnection:
        raise NotImplementedError("子类必须实现create_connection方法")

    def create_processor(self) -> DataProcessor:
        raise NotImplementedError("子类必须实现create_processor方法")


# -------------------------- 6. 具体工厂实现 --------------------------
class TcpFactory(AbstractNetworkFactory):
    def create_connection(self, **kwargs) -> NetworkConnection:
        if "host" not in kwargs or "port" not in kwargs:
            raise ValueError("TCP工厂需要host和port参数")
        return TcpConnection(host=kwargs["host"], port=kwargs["port"])

    def create_processor(self) -> DataProcessor:
        return TcpDataProcessor()


class UdpFactory(AbstractNetworkFactory):
    def create_connection(self, **kwargs) -> NetworkConnection:
        if "host" not in kwargs or "port" not in kwargs:
            raise ValueError("UDP工厂需要host和port参数")
        return UdpConnection(host=kwargs["host"], port=kwargs["port"])

    def create_processor(self) -> DataProcessor:
        return UdpDataProcessor()


class WsFactory(AbstractNetworkFactory):
    def create_connection(self, **kwargs) -> NetworkConnection:
        if "url" not in kwargs:
            raise ValueError("WebSocket工厂需要url参数")
        return WsConnection(url=kwargs["url"])

    def create_processor(self) -> DataProcessor:
        return WsDataProcessor()


# -------------------------- 7. 客户端使用 --------------------------
def use_network(factory: AbstractNetworkFactory, **kwargs):
    """通用网络使用逻辑（依赖抽象，不依赖具体实现）"""
    # 创建产品族：连接 + 处理器
    connection = factory.create_connection(**kwargs)
    processor = factory.create_processor()

    # 业务逻辑：连接 -> 编码数据 -> 发送 -> 关闭
    if connection.connect():
        raw_data = b"Hello, Abstract Factory!"
        encoded_data = processor.encode(raw_data)
        print(f"编码后数据：{encoded_data.hex()}")

        connection.send_data(encoded_data)
        connection.close()


if __name__ == "__main__":
    # 使用TCP产品族
    print("=== TCP 测试 ===")
    tcp_factory = TcpFactory()
    use_network(tcp_factory, host="127.0.0.1", port=8080)

    print("\n=== UDP 测试 ===")
    udp_factory = UdpFactory()
    use_network(udp_factory, host="127.0.0.1", port=8081)

    print("\n=== WebSocket 测试 ===")
    ws_factory = WsFactory()
    use_network(ws_factory, url="ws://127.0.0.1:8082")
